package org.example.offical.doc.ai.service;

import dev.langchain4j.agent.tool.P;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.openai.OpenAiChatModelName;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.model.output.structured.Description;
import dev.langchain4j.service.AiServices;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.example.ToolDemo;
import org.example.offical.doc.ChatMemoryDemo;
import org.example.offical.doc.ModelUtils;
import org.jetbrains.annotations.NotNull;

import java.util.List;
import java.util.stream.Stream;

/**
 * @author superMan
 * @since fish_temp_since
 */
public class AiServiceWithMemoryAndToolExecution {
    public static void main(String[] args) {


        Assistant assistantWithMemory = AiServices.builder(Assistant.class)
                .chatLanguageModel(ModelUtils.getOpenAiDemoModel())
                .chatMemoryProvider(memoryId -> TokenWindowChatMemory.builder()
                        .maxTokens(1000, new OpenAiTokenizer(OpenAiChatModelName.GPT_4_O_MINI))
                        .chatMemoryStore(new ChatMemoryDemo.MyChatMemoryStore())
                        .build()
                )
                .tools(new Tools())
                .build();
        // 演示了LLM Function calling的功能（如果使用AIService 的方式，langchain4j会自动帮我们执行方法，没有使用AIService的方式则需要自己编码调用方法）
        System.out.println(assistantWithMemory.chat("请问李四的BMI是多少"));
        // 演示了记忆的能力，记忆就是langchain4j会将之前的谈话内容再次发送给LLM
        System.out.println("你刚才回答了什么问题，准确吗");

    }


    interface Assistant {
        String chat(String userMsg);
    }


    public static class Tools {

        @Tool("获取用户信息，包括姓名、身高、体重、爱好等")
        public static User getUserInfo(@P("用户的名字") String username) {
            return Stream.of(
                            new User("张三", 18, 170, 60, new String[]{"打篮球", "看电影"}),
                            new User("李四", 20, 180, 70, new String[]{"看电影", "打篮球"}),
                            new User("王五", 22, 190, 80, new String[]{"看电影", "打篮球"}),
                            new User("赵六", 24, 150, 50, new String[]{"看电影", "打篮球"})
                    )
                    .filter(x -> x.name.equals(username))
                    .findFirst()
                    .orElse(null);
        }

        @Tool("计算BMI")
        public static double bmi(@P("体重（KG）") double weight, @P("身高（CM）") double height) {
            return weight / (height * height);
        }
    }

    @Data
    @Description("用户")
    @NoArgsConstructor
    @AllArgsConstructor
    public static class User {
        @Description("姓名")
        private String name;
        @Description("年龄")
        private int age;
        @Description("身高 cm")
        private double height;
        @Description("体重 kg")
        private double weight;
        @Description("爱好")
        private String[] hobby;
    }
}
